package com.welson.db.cryptor.strategy;

import com.welson.db.cryptor.config.DbCryptorProperties;
import org.springframework.stereotype.Component;

import javax.crypto.Cipher;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.IvParameterSpec;

/**
 * AES加解密策略实现
 * @author welson
 */
@Component
public class AESCryptStrategy extends AbstractCryptStrategy {
    private static final String ALGORITHM = "AES";
    /**
     * GCM认证标签长度，单位：位
     */
    private static final int GCM_TAG_LENGTH = 128;

    public AESCryptStrategy(DbCryptorProperties properties) {
        super(properties, ALGORITHM);
    }

    @Override
    public Cipher getCipher(String transformation) throws Exception {
        Cipher cipher = Cipher.getInstance(transformation);
        String mode = properties.getMode().toUpperCase();
        if ("GCM".equals(mode)) {
            // 使用默认的12字节IV长度，但实际使用时可以根据配置的IV长度来初始化
            byte[] iv = properties.getIv() != null ? 
                properties.getIv().getBytes() : new byte[12];
            GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
            cipher.init(Cipher.ENCRYPT_MODE, new javax.crypto.spec.SecretKeySpec(new byte[16], ALGORITHM), gcmParameterSpec);
        }
        return cipher;
    }

    @Override
    public void validateMode(String mode) {
        // AES支持所有模式，无需特殊验证
    }

    @Override
    public boolean isModeRequireIv(String mode) {
        if ("GCM".equals(mode)) {
            return true;
        }
        return super.isModeRequireIv(mode);
    }

    @Override
    public byte[] encrypt(byte[] data, byte[] key, byte[] iv) throws Exception {
        String mode = getCurrentMode();
        String transformation = getTransformation();
        Cipher cipher = getCipher(transformation);
        javax.crypto.spec.SecretKeySpec secretKeySpec = new javax.crypto.spec.SecretKeySpec(key, ALGORITHM);
        
        if ("GCM".equals(mode)) {
            GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
            cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, gcmParameterSpec);
        } else if (isModeRequireIv(mode)) {
            IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
            cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivParameterSpec);
        } else {
            cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec);
        }
        
        return cipher.doFinal(data);
    }

    @Override
    public byte[] decrypt(byte[] data, byte[] key, byte[] iv) throws Exception {
        String mode = getCurrentMode();
        String transformation = getTransformation();
        Cipher cipher = getCipher(transformation);
        javax.crypto.spec.SecretKeySpec secretKeySpec = new javax.crypto.spec.SecretKeySpec(key, ALGORITHM);
        
        if ("GCM".equals(mode)) {
            GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
            cipher.init(Cipher.DECRYPT_MODE, secretKeySpec, gcmParameterSpec);
        } else if (isModeRequireIv(mode)) {
            IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
            cipher.init(Cipher.DECRYPT_MODE, secretKeySpec, ivParameterSpec);
        } else {
            cipher.init(Cipher.DECRYPT_MODE, secretKeySpec);
        }
        
        return cipher.doFinal(data);
    }
} 